#!/usr/bin/python
# -*- coding: UTF-8 -*-

import mongodriver
import json
import sys
import redis

class AnagraphDistance:
    def __init__(self, algorithm_type = 1):
        self.mongo_driver   = mongodriver.MongoDriver()
        self.db             = self.mongo_driver.db['distance']
        self.pool           = redis.ConnectionPool(host='192.168.31.175', port=6379, decode_responses=True)
        self.redis          = redis.Redis(connection_pool=self.pool)

        self.setAlgorithmType(algorithm_type)

    def setAlgorithmType(self, algorithm_type):
        # 只有四中类型算法
        if algorithm_type in [1, 2, 3, 4]:
            self.algorithm_type = algorithm_type

    def getRedisDistance(self, ana_name_src, ana_name_des):
        resKey = ana_name_src + "-" + ana_name_des
        res = self.redis.get(resKey)
        if not res:
            resKey = ana_name_des + "-" + ana_name_src
            res = self.redis.get(resKey)

        return 1.0 - float(res)

    def getDistance(self, ana_name_src, ana_name_des):
        # 欧氏距离
        if self.algorithm_type == 1:
            return self.__findDistance(ana_name_src, ana_name_des, 'euclidean')
        # 余弦相似度
        elif self.algorithm_type == 2:
            return self.__findDistance(ana_name_src, ana_name_des, 'cosine')
        # Jaccard距离
        elif self.algorithm_type == 3:
            return self.__findDistance(ana_name_src, ana_name_des, 'jaccard')
        # person相关系数
        elif self.algorithm_type == 4:
            return self.__findDistance(ana_name_src, ana_name_des, 'person')
        else:
            return None

    def setDistance(self, ana_name_src, ana_name_des, distance):
        # 欧氏距离
        if self.algorithm_type == 1:
            return self.__storeDistance(ana_name_src, ana_name_des, distance, 'euclidean')
        # 余弦相似度
        elif self.algorithm_type == 2:
            return self.__storeDistance(ana_name_src, ana_name_des, distance, 'cosine')
        # Jaccard距离
        elif self.algorithm_type == 3:
            return self.__storeDistance(ana_name_src, ana_name_des, distance, 'jaccard')
        # person相关系数
        elif self.algorithm_type == 4:
            return self.__storeDistance(ana_name_src, ana_name_des, distance, 'person')

    # 删除集合数据
    def flushDistance(self):
        if self.algorithm_type == 1:
            col = self.db['euclidean']
        # 余弦相似度
        elif self.algorithm_type == 2:
            col = self.db['cosine']
        # Jaccard距离
        elif self.algorithm_type == 3:
            col = self.db['jaccard']
        # person相关系数
        elif self.algorithm_type == 4:
            col = self.db['person']

        col.delete_many({})

    def __findDistance(self, ana_name_src, ana_name_des, collection_name):
        col = self.db[collection_name]
        result_dis = col.find_one({'src':ana_name_src, 'des':ana_name_des})
        if result_dis == None:
            result_dis = col.find_one({'src':ana_name_des, 'des':ana_name_src})

        if result_dis == None:
            return None
        else:
            return result_dis['distance']

    def __storeDistance(self, ana_name_src, ana_name_des, distance, collection_name):
        # 将所有距离写入mongoDB
        col = self.db[collection_name]

        if col.find_one({'src':ana_name_src, 'des':ana_name_des}):
            col.find_one_and_update({'src':ana_name_src, 'des':ana_name_des}, {"$set": {"distance": distance}})
        else:
            insert_dict = {}
            insert_dict['src']      = ana_name_src
            insert_dict['des']      = ana_name_des
            insert_dict['distance'] = distance

            col.insert_one(insert_dict)

        if col.find_one({'src':ana_name_des, 'des':ana_name_src}):
            col.find_one_and_update({'src':ana_name_des, 'des':ana_name_src}, {"$set": {"distance": distance}})
        else:
            insert_dict = {}
            insert_dict['src']      = ana_name_des
            insert_dict['des']      = ana_name_src
            insert_dict['distance'] = distance

            col.insert_one(insert_dict)
